import os
import pandas as pd
import time
import argparse
import seaborn as sn
import numpy as np
import matplotlib.pyplot as plt
# sn.set(style="darkgrid") #这是seaborn默认的风格


from model import SASRec
from utils import *


def str2bool(s):
    if s not in {'false', 'true'}:
        raise ValueError('Not a valid boolean string')
    return s == 'true'
def auc_print(x_data,y_data):
    # x_data = pd.DataFrame(x_data)
    # y_data = pd.DataFrame(y_data)
    # objs = [x_data, y_data]
    # data = pd.concat(objs, axis=0, join='outer', ignore_index=False, copy=True)
    ax = sn.lineplot(x=x_data, y=y_data, )
    plt.show()
parser = argparse.ArgumentParser()
'''
基本参数
'''
parser.add_argument('--dataset',default="data/cs-training.csv")# 完整数据集导入、
parser.add_argument('--user_train', default="data/user_train_DT.csv")
parser.add_argument('--user_test', default="data/user_test_DT.csv")
# parser.add_argument('--dataset', default="data/UCI_Credit_Card.csv")
# parser.add_argument('--user_train', default="data/UCI_Credit_Card_train.csv")
# parser.add_argument('--user_test', default="data/UCI_Credit_Card_test.csv")
parser.add_argument('--train_dir', default="cpu")
# parser.add_argument('--train_dir', default="cpu") # 使用cpu训练
parser.add_argument('--batch_size', default=1000, type=int) # 批尺寸
parser.add_argument('--lr', default=0.01, type=float)  # Adam分类器lr学习率
# parser.add_argument('--lr', default=0.01, type=float)
parser.add_argument('--num_epochs', default=30, type=int)
'''
交叉层参数
'''
parser.add_argument('--hidden_units_FM', default=10, type=int)# FM隐向量维度
parser.add_argument('--latent_units', default=10, type=int)
parser.add_argument('--layer_num', default=4, type=int) # CrossNetwork层数 # GMSC最大auc最佳层数6 DCCC最大auc最佳层数4

'''
自注意力层参数
'''
parser.add_argument('--hidden_units', default=10, type=int) # 词向量维度(自注意力层）
parser.add_argument('--num_heads', default=4, type=int) # 多头注意力数量
parser.add_argument('--dropout_rate', default=0.2, type=float) # 丢弃率（dropout）
parser.add_argument('--num_blocks', default=1, type=int) # 自注意力层堆叠
'''
其他参数
'''
parser.add_argument('--l2_emb', default=0.0005, type=float)
parser.add_argument('--device', default='cpu', type=str) # GPU训练
# parser.add_argument('--device', default='cpu', type=str)  # cpu训练
parser.add_argument('--state_dict_path', default=None, type=str)
parser.add_argument('--FMmode', default="CrossNet",
                            type=str)  # ['None','FM','AFM','CrossNet'] # 交叉层切换
parser.add_argument('--Biasmode', default='None', type=str)  # ['lognormal', 'normal', 'abs', 'learn', 'None']
t_auc = []
t_acc = []
t_epoch = []
t_F1 = []
t_auc1 = []

seed = 2022
torch.manual_seed(seed)  # 为CPU设置随机种子
torch.cuda.manual_seed(seed)  # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(seed)  # 为所有GPU设置随机种子
np.random.seed(seed)
random.seed(seed)

args = parser.parse_args()
if not os.path.isdir(args.dataset + '_' + args.train_dir):
    os.makedirs(args.dataset + '_' + args.train_dir)
with open(os.path.join(args.dataset + '_' + args.train_dir, 'args.txt'), 'w') as f:
    f.write('\n'.join([str(k) + ',' + str(v) for k, v in sorted(vars(args).items(), key=lambda x: x[0])]))
f.close()

if __name__ == '__main__':

    dataset = data_partition(args.user_train,args.user_test)

    [user_train, user_test] = dataset
    # emb参数
    emb_list = [len(set(user_train[:,i])) for i in range(1,user_train.shape[1])]
    print(emb_list) #

    parser.add_argument('--emb_list', default=emb_list, type=list)
    parser.add_argument('--maxlen', default=len(emb_list), type=int)
    parser.add_argument('--maxlen_FM', default=sum(emb_list), type=int)
    args = parser.parse_args()

    num_batch = len(user_train) // args.batch_size
    
    f = open(os.path.join(args.dataset + '_' + args.train_dir, 'log.txt'), 'w')
    sampler = WarpSampler(user_train, batch_size=args.batch_size, n_workers=1)

    model = SASRec(args).to(args.device)
    # model-SASRec
    for name, param in model.named_parameters():
        try:
            torch.nn.init.xavier_normal_(param.data)
        except:
            pass


    model.train()


    epoch_start_idx = 1
    if args.state_dict_path is not None:
        try:
            model.load_state_dict(torch.load(args.state_dict_path, map_location=torch.device(args.device)))
            tail = args.state_dict_path[args.state_dict_path.find('epoch=') + 6:]
            epoch_start_idx = int(tail[:tail.find('.')]) + 1
        except: # in case your pytorch version is not 1.6 etc., pls debug by pdb if load weights failed
            print('failed loading state_dicts, pls check file path: ', end="")
            print(args.state_dict_path)
            print('pdb enabled for your quick check, pls type exit() if you do not need it')
            import pdb; pdb.set_trace()

    bce_criterion = torch.nn.BCEWithLogitsLoss().to(args.device)
    adam_optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.98))

    
    T = 0.0
    t0 = time.time()
    
    for epoch in range(epoch_start_idx, args.num_epochs + 1):
        for step in range(num_batch):
            seq, pos = sampler.next_batch()
            seq, pos = np.array(seq), torch.FloatTensor(pos).unsqueeze_(1).to(args.device)
            pre = model(seq)

            adam_optimizer.zero_grad()
            loss = bce_criterion(pre, pos)

            loss.backward()
            adam_optimizer.step()

        if epoch % 1 == 0:
            model.eval()
            t1 = time.time() - t0
            T += t1
            print('Evaluating', end='')
            with torch.no_grad():
                test_auc, test_acc, test_recall, test_ap, test_f1 = evaluate(model, user_test, args)
                train_auc, train_acc, train_recall, train_ap, train_f1 = evaluate(model, user_train, args)
            print('epoch:%d, time: %f(s), valid (AUC: %.4f,ACC: %.4f,Recall: %.4f,AP: %.4f,F1: %.4f),'
                  ' test (AUC: %.4f,ACC: %.4f,Recall: %.4f,AP: %.4f,F1: %.4f)'
                    % (epoch, T, train_auc, train_acc, train_recall, train_ap, train_f1, test_auc, test_acc, test_recall, test_ap, test_f1))

            t_auc.append(test_auc)
            t_acc.append(test_acc)
            t_epoch.append(epoch)
            t_F1.append(test_f1)
    
            f.write(str(train_auc) + ' ' + str(test_auc) + '\n')
            f.flush()
            t0 = time.time()
            model.train()
    
        if epoch == args.num_epochs:
            folder = args.dataset + '_' + args.train_dir
            fname = 'SASRec.epoch={}.lr={}.layer={}.head={}.hidden={}.pth'
            fname = fname.format(args.num_epochs, args.lr, args.num_blocks, args.num_heads, args.hidden_units)
            torch.save(model.state_dict(), os.path.join(folder, fname))

    t_auc1 = t_auc.copy()
    t_auc1.sort()
    t_auc1 = t_auc1[2:-1]
    print('测试集最大AUC {:.6f} 测试集平均AUC {:.6f} 测试集最大F1 {:.6f} 测试集平均F1 {:.6f}'.format(float(pd.DataFrame(t_auc).max()),
    float(pd.DataFrame(t_auc1).mean()), float(pd.DataFrame(t_F1).max()), float(pd.DataFrame(t_F1).mean())))

    f.close()
    sampler.close()
    print(str(args))
    print("Done")
    auc_print(t_epoch, t_auc)

